import os
import sys
from html.parser import HTMLParser
from io import StringIO

from tqdm import tqdm

from zkl_datasets import ChunkedBinaryDatasetWriter, load_dataset

root_dir_path = os.path.join(os.path.dirname(__file__), "../../..")
sys.path.append(root_dir_path)

from scripts.config import datasets_dir_path

org_dataset_path = os.path.join(datasets_dir_path, "DuReader")
dst_dataset_path = os.path.join(datasets_dir_path, "DuReader-cleaned")


class MyHTMLParser(HTMLParser):
    def __init__(self, io: StringIO):
        super().__init__()
        self.io = io
        self.stack = []

    def handle_starttag(self, tag, attrs):
        self.stack.append(tag)

        if tag in ['p', 'br', 'div', 'table', 'tr', 'hr']:
            self.io.write("\n")
        elif tag in ['img']:
            self.io.write("[图片]")
        elif tag in ['audio', 'ikaudio']:
            self.io.write("[音频]")
        elif tag in ['video', 'ikvideo']:
            self.io.write("[音频]")
        elif tag in ['file']:
            self.io.write("[文件]")
        elif tag in ['iframe', 'embed']:
            self.io.write("[嵌入内容]")
        elif tag in ['strong', 'sup', 'sub', 'pre', 'u', 'em', 'i', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
            pass
        elif tag in ['a', 'ul', 'ol', 'li', 'td', 'blockquote', 'form', 'tbody']:
            pass
        else:
            raise ValueError(f"Unknown tag: {tag}")

    def handle_endtag(self, tag):
        popped_tag = self.stack.pop()
        assert popped_tag == tag

    def handle_data(self, data):
        self.io.write(data)


def process_paragraph(paragraph: str) -> str:
    if paragraph.startswith('<') and paragraph.endswith('>'):
        try:
            io = StringIO()
            parser = MyHTMLParser(io)
            parser.feed(paragraph)
            return io.getvalue()
        except Exception:
            return paragraph
    return paragraph


if __name__ == '__main__':
    org_dataset = load_dataset(org_dataset_path, split_into_documents=True)
    dst_dataset_writer = ChunkedBinaryDatasetWriter(dst_dataset_path)
    for sample in tqdm(org_dataset):
        paragraphs = [process_paragraph(paragraph) for paragraph in sample['paragraphs']]
        text = sample['title'] + "\n" + '\n'.join(paragraphs)
        dst_dataset_writer.append(text.encode('utf-8'))
